import keyboard
import math
import cmath
import time
import random
import numpy as np
import matplotlib.pyplot as plt


class Sequence:
    num = 0

    def __init__(self, start=0, end=0, value=0, infinite=False, /, *, from_value=False, value_list=None):
        self.dft = None
        self.sequence = {}
        self.start = start
        self.end = end
        self.infinite = infinite
        if self.infinite and not from_value:
            def esc_callback():
                self.infinite = False
                print("\n停止输入无限长序列, 按任意键继续...")

            keyboard.add_hotkey('esc', esc_callback)
            while self.infinite:
                val_input = input("请输入下一个序列值(或esc退出):")
                # print(val_input)
                if val_input.isdigit():
                    self.sequence[start] = int(val_input)  # 假定所有值都是整数，否则改为float
                    start += 1
            self.end = start - 1
            # keyboard.wait('esc')
            keyboard.remove_hotkey('esc')

        elif not from_value:
            self.sequence = {i: value for i in range(start, end + 1)}
        else:
            self.sequence = {i: value_list[i - start] for i in range(start, start + len(value_list))}
            self.end = self.start + len(value_list) - 1
        self.id = Sequence.num
        Sequence.num += 1

    def __iter__(self):
        return iter(self.sequence)

    def __len__(self):
        return len(self.sequence)

    def __getitem__(self, key):
        return self.sequence.get(key)

    def __setitem__(self, key, value):
        self.sequence[key] = value

    def print(self):
        values = list(self.sequence.values())
        for i in range(0, len(values)):
            values[i] = complex(round(values[i].real, 4), round(values[i].imag, 4))
        print('序列{} = '.format(self.id), values, ', n={}:{}'.format(self.start, self.end),
              sep='')

    def dft_print(self):
        if self.dft is not None:
            values = list(self.dft.values())
            for i in range(0, len(values)):
                values[i] = complex(round(values[i].real, 4), round(values[i].imag, 4))
            print('DFT(序列{}) = '.format(self.id), values, ', n={}:{}'.format(self.start, self.end),
                  sep='')
        else:
            print('序列{}尚未计算DFT!'.format(self.id))

    def padding(self, left, right):
        # temp_dict = self.sequence.copy()
        # for i in range(self.start - left, self.start):
        #     temp_dict[i] = 0
        # for i in range(self.end +1, self.end+right+1):
        #     temp_dict[i] = 0
        # return dict(sorted(temp_dict.items()))
        for i in range(self.start - left, self.start):
            self.sequence[i] = 0
        for i in range(self.end + 1, self.end + right + 1):
            self.sequence[i] = 0
        self.start = self.start - left
        self.end = self.end + right
        self.sequence = dict(sorted(self.sequence.items()))

    def shift(self, value):
        self.start = self.start - value
        self.end = self.end - value
        temp_sequence = self.sequence.copy()
        self.sequence.clear()
        for i in range(self.start, self.end + 1):
            self.sequence[i] = temp_sequence[i + value]
        temp_sequence.clear()

    def reverse(self):
        temp_sequence = self.sequence.copy()
        self.sequence.clear()
        self.start, self.end = -self.end, -self.start
        for i in range(self.start, self.end + 1):
            self.sequence[i] = temp_sequence[-i]
        temp_sequence.clear()

    def downsample(self, factor):
        temp_sequence = self.sequence.copy()
        self.sequence.clear()
        self.start = int(self.start / factor)
        self.end = int(self.end / factor)
        for i in range(self.start, self.end + 1):
            self.sequence[i] = temp_sequence[i * factor]
        temp_sequence.clear()

    def upsample(self, factor):
        temp_sequence = self.sequence.copy()
        self.sequence.clear()
        self.start = self.start * factor
        self.end = self.end * factor
        for i in range(self.start, self.end + 1):
            if i % factor == 0:
                self.sequence[i] = temp_sequence[i // factor]
            else:
                self.sequence[i] = 0
        temp_sequence.clear()

    def difference(self):
        temp_sequence = self.sequence.copy()
        self.sequence.clear()
        self.end = self.end - 1
        for i in range(self.start, self.end + 1):
            self.sequence[i] = temp_sequence[i + 1] - temp_sequence[i]
        temp_sequence.clear()

    def sum(self, left, right):
        left = max(left, self.start)
        right = min(right, self.end)
        return sum([self.sequence[i] for i in range(left, right + 1)])

    def calc_dft(self):
        pi = math.pi
        e = math.e
        N = self.end - self.start + 1
        self.dft = {}
        for i in range(self.start, self.end + 1):
            self.dft[i] = 0
            for k in range(self.start, self.end + 1):
                self.dft[i] += self.sequence[k] * (e ** (-1j * 2 * pi / N * i * k))

    def calc_idft(self):
        pi = math.pi
        e = math.e
        N = self.end - self.start + 1
        idft = {}
        for i in range(self.start, self.end + 1):
            idft[i] = 0
            for k in range(self.start, self.end + 1):
                idft[i] += self.dft[k] * (e ** (1j * 2 * pi / N * i * k))
            idft[i] /= N
        return idft


def add_sequences(a, b):
    left = min(a.start, b.start)
    right = max(a.end, b.end)
    a.padding(a.start - left, right - a.end)
    b.padding(b.start - left, right - b.end)
    values_a = a.sequence.values()
    values_b = b.sequence.values()
    return Sequence(left, from_value=True, value_list=[x + y for x, y in zip(values_a, values_b)])


def mul_sequences(a, b):
    left = min(a.start, b.start)
    right = max(a.end, b.end)
    a.padding(a.start - left, right - a.end)
    b.padding(b.start - left, right - b.end)
    values_a = a.sequence.values()
    values_b = b.sequence.values()
    return Sequence(left, from_value=True, value_list=[x * y for x, y in zip(values_a, values_b)])


def conv_sequences(a, b):
    len_a = len(a.sequence)
    len_b = len(b.sequence)
    list_a = list(a.sequence.values())
    list_b = list(b.sequence.values())

    list_b.reverse()
    list_a[:0] = [0 for i in range(len_b - 1)]
    list_a = list_a + [0 for i in range(len_b - 1)]
    result = []
    for i in range(len_a + len_b - 1):
        result.append(sum([x * y for x, y in zip(list_a[i:], list_b)]))
    return Sequence(a.start + b.start, from_value=True, value_list=result)


def window_similarity(a, b, lag):
    b.shift(-lag)
    total = 0
    for i in range(b.start, b.end + 1):
        value_a = a[i] if a[i] is not None else 0
        total = total + value_a * b[i]
    return total


def normalized_similarity(a, b):
    left = min(a.start, b.start)
    right = max(a.end, b.end)
    a.padding(a.start - left, right - a.end)
    b.padding(b.start - left, right - b.end)
    values_a = a.sequence.values()
    values_b = b.sequence.values()
    mean_a = sum(values_a) / len(values_a)
    mean_b = sum(values_b) / len(values_b)
    zero_values_a = [value - mean_a for value in values_a]
    zero_values_b = [value - mean_b for value in values_b]
    temp_a = sum(value ** 2 for value in zero_values_a) ** 0.5
    temp_b = sum(value ** 2 for value in zero_values_b) ** 0.5
    normalized_a = [value / temp_a for value in zero_values_a]
    normalized_b = [value / temp_b for value in zero_values_b]
    return sum(a * b for a, b in zip(normalized_a, normalized_b))


def main():
    # t = Sequence(0, from_value=True, value_list=[2, 4, 6, 7, 11, 12, 13, 14])
    # s.print()
    # t.print()
    # result = add_sequences(s, t)
    # result = mul_sequences(s, t)
    # result = conv_sequences(s, t)
    # result = window_similarity(s, t, -1)
    # result = normalized_similarity(s, t)
    # result.print()
    # print(result)
    # s.print()
    # t.print()

    # s.dft_print()
    # s.calc_dft()
    # s.dft_print()
    # print(s.calc_idft())

    # 比较实现的DFT和FFT的时间
    # length = 1000
    # random_list1 = [random.randint(1, 1000) for _ in range(length)]
    # random_list2 = random_list1.copy()
    #
    # s = Sequence(0, from_value=True, value_list=random_list1)
    # start_time = time.perf_counter()
    # s.calc_dft()
    # end_time = time.perf_counter()
    # print(f"CPU时间：{end_time - start_time:.5f}秒")
    #
    # start_time = time.perf_counter()
    # np.fft.fft(random_list2)
    # end_time = time.perf_counter()
    # print(f"CPU时间：{end_time - start_time:.5f}秒")

    # 验证频域乘积等于时域卷积
    # pi = math.pi
    # length = 20
    # x = [math.cos(0.2 * pi * i) + math.cos(0.4 * pi * i) + math.cos(0.6 * pi * i) for i in range(length)]
    # y = [math.cos(0.4 * pi * i) for i in range(length)]
    # s = Sequence(0, from_value=True, value_list=x)
    # t = Sequence(0, from_value=True, value_list=y)
    # s.calc_dft()
    # t.calc_dft()
    # # 卷积需要周期延拓，手动实现
    # r_list = []
    # for i in range(length):
    #     total = 0
    #     for j in range(length):
    #         total += x[j] * y[(i - j + length) % length]
    #     r_list.append(total)
    # r = Sequence(0, from_value=True, value_list=r_list)
    # r.calc_dft()
    # s.dft_print()
    # t.dft_print()
    # r.dft_print()
    #
    # s.print()
    # t.print()
    # r.print()

    # 绘制频谱图
    pi = math.pi
    length = 20
    zeros = 0
    x = [3 * math.cos(0.2 * pi * i) + 4 * math.sin(0.2 * pi * i) + 6 * math.cos(0.7 * pi * i) + 2 * math.cos(
        0.9 * pi * i) for i in range(length)]
    # x = [random.randint(-100, 100) for _ in range(length)]
    for i in range(zeros):
        x.append(0)
    length += zeros
    s = Sequence(0, from_value=True, value_list=x)
    s.calc_dft()
    s.dft_print()

    dft_list = list(s.dft.values())
    frequency = [1 / length * i for i in range(length)]
    amplitude = [math.sqrt(dft_list[i].real ** 2 + dft_list[i].imag ** 2) / length for i in range(length)]
    phase = [cmath.phase(dft_list[i]) for i in range(length)]

    plt.subplot(2, 1, 1)
    # plt.plot(frequency,amplitude)
    plt.scatter(frequency, amplitude, marker='o')
    plt.grid(False)
    plt.xlabel('Frequency in π units')
    plt.ylabel('Amplitude')
    plt.ylim(0)
    for i in range(length):
        plt.vlines(frequency[i], 0, amplitude[i])

    plt.subplot(2, 1, 2)
    plt.plot(frequency, phase)
    plt.xlabel('Frequency in π units')
    plt.ylabel('Phase')

    plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    main()
